#!/usr/bin/env python3

"""
Base Camera Class for Image and World Position Transformations
所有的相机类都继承自这个类, 相机类中实现成员变量
D435Camera, D435iCamera
"""

import cv2
import numpy as np
from typing import Optional

class BaseCamera:
    width: int
    height: int
    only_rgb: bool
    K: np.ndarray  # Camera intrinsic matrix
    dist: np.ndarray  # Distortion coefficients

    def __init__(
            self, rgb_width: int, rgb_height: int, fps: int, only_rgb: bool,
            depth_width: Optional[int] = None, depth_height: Optional[int] = None,
            **kwargs
        ):
        self.width, self.height, self.fps, self.only_rgb, self.depth_width, self.depth_height = (
            rgb_width, rgb_height, fps, only_rgb, depth_width, depth_height
        )

    def get_frame(self) -> tuple[np.ndarray, np.ndarray]:
        """ Get RGB and Depth frame from camera
        Returns:
            rgb_img (shape=(H,W,3)): RGB image
            depth_img (shape=(H,W)): Depth image in millimeters
        """
        raise NotImplementedError("This method should be implemented by subclasses.")

    def world_pos2img_pos(self, xs_world: np.ndarray | list, rotation: np.ndarray, translation: np.ndarray):
        """ Convert world positions to image position
        Args:
            xs_world (shape=(N,3)): World positions
            rotation (shape=(3,3)): Rotation matrix (World to Camera coord-sys)
            translation (shape=(3,)): Translation matrix (World to Camera coord-sys)
        Returns:
            xs_img (shape=(N,2)): Width and Height
        """
        if not isinstance(xs_world, np.ndarray):
            xs_world = np.array(xs_world)
        if xs_world.ndim == 1:
            xs_world = xs_world.reshape(1, -1)

        T = translation.reshape(3, 1)
        R = rotation
        xs_camera = xs_world @ R.T + T.T
        xs_img = xs_camera @ self.K.T / xs_camera[:, -1:]
        xs_img = xs_img[:, :2].astype(np.int32)
        return xs_img

    def img_pos2camera_pos(self, xs_pos: np.ndarray | list, depth_img: Optional[np.ndarray] = None, xs_depth: Optional[np.ndarray | list] = None):
        """ Convert image positions to camera position
        Args:
            xs_pos (shape=(2,) or (N,2)): Image positions (W, H)
            depth_img (shape=(H,W)): Depth image, or give xs_depth
            xs_depth (shape=(N,)): Depth info for each img pos, or give depth_img
        """
        assert depth_img is not None or xs_depth is not None
        if not isinstance(xs_pos, np.ndarray):
            xs_pos = np.array(xs_pos)
        if xs_depth is not None and not isinstance(xs_depth, np.ndarray):
            xs_depth = np.array(xs_depth)
        xs_pos = xs_pos.copy()
        if xs_pos.ndim == 1:
            xs_pos = xs_pos.reshape(1, -1)
        N = xs_pos.shape[0]
        xs_pos[:, 0] = np.clip(xs_pos[:, 0], 0, self.width-1)
        xs_pos[:, 1] = np.clip(xs_pos[:, 1], 0, self.height-1)
        xs_pos = np.concatenate([xs_pos, np.ones((N, 1), xs_pos.dtype)], axis=-1)

        # Image to Camera coord-sys
        K_inv = np.linalg.inv(self.K)
        xs = xs_pos @ K_inv.T
        r2 = xs[:, 0] ** 2 + xs[:, 1] ** 2
        k1, k2, p1, p2 = self.dist
        x, y = xs[:, 0], xs[:, 1]
        xs[:, 0] = x * (1+k1*r2+k2*r2**2) + 2*p1*x*y + p2*(r2+2*x**2)
        xs[:, 1] = y * (1+k1*r2+k2*r2**2) + p1*(r2+2*y**2) + 2*p2*x*y

        xs_img_idxs = xs_pos.astype(np.int32)
        if xs_depth is None:
            xs_depth = depth_img[xs_img_idxs[:,1], xs_img_idxs[:,0]] / 1e3  # NOTE: millimeter -> meter
        xs_camera = xs * xs_depth.reshape(-1, 1)

        return xs_camera

    def img_pos2world_pos(
            self, xs_pos: np.ndarray | list,
            rotation: np.ndarray, translation: np.ndarray,
            depth_img: Optional[np.ndarray] = None, xs_depth: Optional[np.ndarray | list] = None,
        ):
        """ Convert image positions to world position
        Args:
            xs_pos (shape=(2,) or (N,2)): Image positions
            rotation (shape=(3,3)): Rotation matrix (World to Camera coord-sys)
            translation (shape=(3,)): Translation matrix (World to Camera coord-sys)
            depth_img (shape=(H,W)): Depth image, or give xs_depth
            xs_depth (shape=(N,)): Depth for each img pos, or give depth_img
        """
        xs_camera = self.img_pos2camera_pos(xs_pos, depth_img, xs_depth)

        # Camera to World corrd-sys
        T = translation.reshape(3, 1)
        R = rotation
        xs_world = (xs_camera - T.T) @ R  # NOTE: R.inv.T == R
        return xs_world
    
    def draw_axis(
            self,
            img: np.ndarray, rotation: np.ndarray, translation: np.ndarray,
            axis_size: float = 0.1,
            axis_colors = ((0,0,255), (0,255,0), (255,0,0)),
            center_color = (255,255,255)):
        """ Draw world translation in image, inplace
        Args:
            img (shape=(H,W,C)): Draw img
            rotation (shape=(3,3)): Rotation matrix
            translation (shape=(3,)): Translation vector
            axis_size: Axis draw size, meter
            axis_colors: Colors of x, y, z axis
            center_color: Color of the center point
        """
        axis_world_pos = np.array([[0,0,0], [axis_size,0,0], [0,axis_size,0], [0,0,axis_size]], np.float32)
        axis_img_pos = self.world_pos2img_pos(axis_world_pos, rotation, translation)
        origin_img_pos = axis_img_pos[0]
        for idx, color in enumerate(axis_colors):
            img = cv2.line(img, origin_img_pos, axis_img_pos[idx+1], color, thickness=2)
        img = cv2.circle(img, origin_img_pos, radius=5, color=center_color, thickness=-1)
        return img
